import cv2
import time
import ulid
from enum import Enum

from config import config
from utils import linear


class ObjectTracer:
    TRACE_TIME_THRESHHOLD = linear(0.5, 2, 1 - config.alert_sensitivity)  # 确认一次物体状态所需的时间（秒）
    TRACE_ENTER_PERCENT_THRESHHOLD = linear(0.4, 0.9, 1 - config.alert_sensitivity)  # 确认物体进入时所需的检测比（从无到有）
    TRACE_LEAVE_PERCENT_THRESHHOLD = 0.  # 确认物体离开时所需的检测比（从有到无）
    TRACE_MIN_SIM = 0.5  # 最小相似度，高于这个相似度的可以认为是之前就在画面中的物体
    DISABLE_DRAW_ALERT_BOX = not config.enable_draw_alert_box  # 报警图像是否禁用绘制预测框

    class Action(Enum):
        NONE = 0
        ENTER = 1
        LEAVE = 2
        FAKE = 3  # 当之前不存在，并且检测之后还是不存在时触发，不用再追踪该对象了

    def __init__(self, label, box, image):
        self.id = ulid.new().str
        self.label = label
        self.box = box
        self.enter_pos = self.get_pos()
        self.create_time = time.perf_counter()
        self.current_time = self.create_time
        self.image = image
        self.exists = False
        self._last_trace_time = self.create_time  # 上一个检测的时间点
        self._ex_trace_count = 1  # 记录存在的帧数
        self._total_trace_count = 1  # 记录检测的总帧数

    def sim(self, box):
        xmin1, ymin1, xmax1, ymax1 = self.box
        xmin2, ymin2, xmax2, ymax2 = box
        w1, h1, w2, h2 = xmax1 - xmin1, ymax1 - ymin1, xmax2 - xmin2, ymax2 - ymin2

        inter = max(0, min(xmax1, xmax2) - max(xmin1, xmin2)) * max(0, min(ymax1, ymax2) - max(ymin1, ymin2))
        union = w1 * h1 + w2 * h2 - inter
        if inter <= 0 or union <= 0:
            iou = 0
        else:
            iou = inter / union

        xmin = min(xmin1, xmin2)
        ymin = min(ymin1, ymin2)
        xmax = max(xmax1, xmax2)
        ymax = max(ymax1, ymax2)
        w, h = xmax - xmin, ymax - ymin

        # 中心点距离相似度（0~1）
        try:
            dis_sim = 1 - (((xmin1 + xmax1) / 2 - (xmin2 + xmax2) / 2) ** 2 + ((ymin1 + ymax1) / 2 - (ymin2 + ymax2) / 2) ** 2) / (w ** 2 + h ** 2)
        except ZeroDivisionError:
            dis_sim = 1

        # 形状相似度
        try:
            shape_sim = 1 - ((w1 - w2) ** 2 / w ** 2 + (h1 - h2) ** 2 / h ** 2) / 2
        except ZeroDivisionError:
            shape_sim = 1
        return iou * 0.4 + dis_sim * 0.2 + shape_sim * 0.4
    
    def update(self, time, box=None, image=None):
        # 当box为None时，表示该物体不在，空更新一次
        self.current_time = time
        if box is not None:
            self._ex_trace_count += 1
            self.box = box
            self.image = image
        self._total_trace_count += 1

        if time - self._last_trace_time > self.TRACE_TIME_THRESHHOLD:  # 当经过一次追踪检测之后，根据当前状态更新动作
            self._last_trace_time = time
            # 计算检测比
            trace_percent = self._ex_trace_count / self._total_trace_count if self._total_trace_count else 0
            # 给计数器置零
            self._ex_trace_count = self._total_trace_count = 0

            if not self.exists:
                if trace_percent >= self.TRACE_ENTER_PERCENT_THRESHHOLD:
                    # 判定进入时，需要达到进入检测比阈值
                    self.exists = True
                    return self.Action.ENTER
                else:
                    return self.Action.FAKE
            elif trace_percent <= self.TRACE_LEAVE_PERCENT_THRESHHOLD:  # 能走这个elif，肯定是existed=True
                # 判定离开时，需要小于离开检测比阈值，离开对检测比的判定稍微宽松点，因为即使只有一次检测到了，也能证明物体存在
                return self.Action.LEAVE
        return self.Action.NONE
    
    # 给这个物体截下来，只有当该物体存在或者进入时，才会设置图片，如果时离开或者Fake，那就不设置
    def make_shotcut(self, image=None):
        if self.DISABLE_DRAW_ALERT_BOX:
            return self.image
        if image is None:
            image = self.image.copy()
        cv2.rectangle(image, self.box[:2], self.box[2:4], (247, 200, 125), 2)
        return image
    
    def get_pos(self):
        return ((self.box[0] + self.box[2]) / 2, (self.box[1] + self.box[3]) / 2)
